import json
from pathlib import Path
from typing import Any,Callable,Dict,List,Optional,Union

from langchain_core.documents import Document
from langchain_community.document_loaders.base import BaseLoader
import re

class CustomJsonLoader(BaseLoader):
    def __init__(self,
                 filePath:Union[str,Path],
                 jqSchema:str):
        import jq
        self.filePath = Path(filePath).resolve()
        self._jqSchema = jq.compile(jqSchema)

    def load(self) -> list[Document]:
        """从 JSON 文件中加载并返回文档。"""
        docs:List[Document] = []
        # print(self.filePath.read_text(encoding="utf-8"))
        self._parse(self.filePath.read_text(encoding="utf-8"),docs)
        return docs

    def _parse(self,content:str,docs:List[Document]) -> None:
        """将给定内容转换为文档。"""
        # content : 原始的JSON
        # json.loads(content)：JSON 格式的字符串 content 转换为 Python 的数据结构。具体转换为哪种形式的数据结构取决于原始 JSON 字符串的内容：

        # 根据jq编译查找到的结果
        # print(content)
        data = self._jqSchema.input(json.loads(content)).all()
        print(self._jqSchema)
       # print(json.loads(content))
        print(self._jqSchema.input(json.loads(content)))
        print(self._jqSchema.input(json.loads(content)).all())
        # print(data)
        for i,sample in enumerate(data,len(docs)-1):
            meta_data={"result":"\n".join(sample), "source": self.filePath}
            docs.append(Document(page_content="\n".join(sample), metadata=meta_data))


loader = CustomJsonLoader(filePath="f:\\cn_test_set.json",
                          jqSchema='."Session-3"."对话历史"')
data = loader.load()
print(data)